"港航 AI 识别完毕后处理代码"

__author__ = "zwj"

# 导入根目录作为模块搜索路径
import os, sys

# ([".."] * 4) 就是项目目录里面 4 层
parent_path = os.path.abspath(os.path.join(__file__, *([".."] * 4)))
sys.path.insert(0, parent_path)

import yaml, time, cv2
from typing import Dict, List
from enum import Enum, unique
import math
from collections import deque

from core.task.modules.processors.base_processor import BaseProcessor
from core.utils.state_machine import StateMachine
from core.utils.logger import logging

from projects.GangHang.processors.esdk import esdk
from projects.GangHang.processors.http_message import HttpMessage, CameraType


# 船只航向枚举
@unique
class Heading(Enum):
    kLeft = "left"
    kRight = "right"
    kUnknown = "unknown"


# 像素点坐标
class PointData:
    def __init__(self, x=0, y=0):
        # 确保点坐标都为正数
        self.x = abs(x)
        self.y = abs(y)


# 船只检测数据汇总
class ShipData:
    def __init__(
        self,
        track_id,
        ship_start: PointData,
        ship_end: PointData,
    ):
        # 船只检测
        self.track_id = track_id
        self.ship_start = ship_start
        self.ship_end = ship_end

    def get_midpoint_data(self) -> PointData:
        return PointData(
            (self.ship_start.x + self.ship_end.x) / 2,
            (self.ship_start.y + self.ship_end.y) / 2,
        )

    def get_area(self) -> int:
        return (self.ship_end.x - self.ship_start.x) * (
            self.ship_end.y - self.ship_start.y
        )


# 船只检测数据
class Mask:
    def __init__(self, width, height):
        self.__width = width
        self.__height = height
        # 初始化一个二维列表
        self.clear()

    def set_mask(self, start_point: PointData, end_point: PointData):
        # 确保数据不过界
        start_point = self.__point_data_int(start_point)
        end_point = self.__point_data_int(end_point)

        # 将给定框范围内的元素置为True
        for y in range(start_point.y, end_point.y):
            for x in range(start_point.x, end_point.x):
                self.__mask[y][x] = True

    def get_mask(self, point: PointData):
        point = self.__point_data_int(point)
        return self.__mask[int(point.y)][int(point.x)]

    def clear(self):
        self.__mask = [
            [False for _ in range(self.__width)] for _ in range(self.__height)
        ]

    def __point_data_int(self, point: PointData):
        point_x = min(int(point.x), self.__width - 1)
        point_y = min(int(point.y), self.__height - 1)
        return PointData(int(point_x), int(point_y))


@unique
class StateMachineEnum(Enum):
    kDetectReady = 0  # 检测准备
    kDetect = 1  # 检测船只,锁定要处理的船只
    kZoom = 2  # 变焦
    kZoomDetect = 3  # 变焦后检测
    kCameraReset = 4  # 镜头复位


class GangHangProcessor(BaseProcessor):
    def __init__(
        self,
        keys,
        Custom_cls_map,
        Screen_size,
        Left_Heading,
        Test,
        HttpMessage_cfg_path,
    ) -> None:

        # * key 值对应索引
        self.__keys = keys

        # * 输入数据
        self.__in_key = self.__keys["in"]
        self.__track_key = self.__keys["track"]
        self.__out_key = self.__keys["out"]
        self.__finish_tag_key = self.__keys["finish_tag"]

        # * 检测 id 对应含义地图
        if isinstance(Custom_cls_map, str):
            with open(Custom_cls_map, "r") as f:
                self.__cls_map = {
                    x: y for x, y in enumerate(yaml.safe_load(f)["label_list"])
                }
        else:
            self.__cls_map = Custom_cls_map

        # * 绘图框颜色初始化, BGR 颜色, 0 0 0 为黑色
        self.__color = (0, 0, 255)

        # * 图片输出目录初始化
        self.__image_output_path = os.getcwd() + "/output"
        if not os.path.exists(self.__image_output_path):
            os.makedirs(self.__image_output_path)

        # * 获取画面大小
        self.__screen_size_x = Screen_size[0]
        self.__screen_size_y = Screen_size[1]

        # * 状态机
        self.__state_machine = StateMachine()
        self.__state_machine.set_state_config(
            StateMachineEnum.kDetectReady, self.__detect_ready
        )
        self.__state_machine.set_state_config(StateMachineEnum.kDetect, self.__detect)
        self.__state_machine.set_state_config(StateMachineEnum.kZoom, self.__zoom)
        self.__state_machine.set_state_config(
            StateMachineEnum.kZoomDetect, self.__zoom_detect
        )
        self.__state_machine.set_state_config(
            StateMachineEnum.kCameraReset, self.__camera_reset
        )
        self.__state_machine.set_initial_state(StateMachineEnum.kDetectReady)

        # * 船只数据
        self.__ship_data_list: List[ShipData] = []

        self.__detect_count = 0
        self.__detect_threshold = 5

        self.__zoom_sleep_count = 0
        self.__zoom_sleep_threshold = 30
        self.__zoom_point = PointData(0, 0)
        self.__zoom_max_area = 0

        self.__ship_processing_map = {}
        self.__position_offset_map = {}
        # 移动阈值以画幅比例作为依据
        self.__position_offset_threshold = 1.0

        # 船只航向要求
        self.__set_ship_heading(Left_Heading)

        # 船只掩码
        self.__mask = Mask(self.__screen_size_x, self.__screen_size_y)
        # 初始化一个掩码二维列表,将不检测区屏蔽
        self.__mask.set_mask(
            PointData(0, 0),
            PointData(self.__screen_size_x, self.__screen_size_y // 3),
        )
        if self.__ship_heading == Heading.kLeft:
            self.__mask.set_mask(
                PointData(0, 0),
                PointData(self.__screen_size_x // 2, self.__screen_size_y),
            )
        elif self.__ship_heading == Heading.kRight:
            self.__mask.set_mask(
                PointData(self.__screen_size_x // 2, 0),
                PointData(self.__screen_size_x, self.__screen_size_y),
            )

        self.__mask_list = []
        self.__zoom_list = []

        # 运动船只检测
        self.__sports_ship_list: List[ShipData] = {}

        # * 控制接口初始化
        self.__http_message = HttpMessage()
        self.__http_message.init_cfg(HttpMessage_cfg_path)

        # * DEBUG : 测试整个追踪结果
        self.__debug = False

        self.__test = Test

        # * 开启机场无人机直播流
        false_count = 0
        if not self.__test:
            self.__http_message.live_stop()
        logging.info("开启无人机直播流...")
        while True:
            if self.__test:
                break
            if esdk.StartH264Stream():
                break
            else:
                false_count += 1
                self.__http_message.live_stop()
                if false_count > 5:
                    logging.error("无人机直播流开启失败!")
                    self.close()
        logging.info("无人机直播流开启完毕!")

        # * 记录操作时间戳
        self.__last_fps_time = time.time()

    def __set_ship_heading(self, Left_Heading):
        if Left_Heading:
            self.__ship_heading = Heading.kLeft
        else:
            self.__ship_heading = Heading.kRight

    def __refresh_ship_data(self, data: Dict) -> Dict:
        # 获取结果
        track_result = data[self.__track_key]

        ship_data_list = []

        for bboxes, track_ids in zip(track_result["boxes"], track_result["track_ids"]):

            for box, track_id in zip(bboxes, track_ids):
                cls_id, score, x1, y1, x2, y2 = box
                cls_id = int(cls_id)

                # 只记录检测区域内船只
                ship_data = ShipData(track_id, PointData(x1, y1), PointData(x2, y2))

                midpoint_data = ship_data.get_midpoint_data()

                # 获取到的船中心点如果在掩码区域内跳过
                if self.__mask.get_mask(midpoint_data):
                    continue

                ship_data_list.append(ship_data)

        self.__ship_data_list = ship_data_list

        return data

    def __detect_image(self, data: Dict) -> Dict:
        image = data[self.__in_key][0].copy()
        draw_thickness = min(image.shape[:2]) // 320

        for start, end in self.__zoom_list:
            # 将变焦框标注
            image = cv2.rectangle(
                image,
                (int(start.x), int(start.y)),
                (int(end.x), int(end.y)),
                color=self.__color,
                thickness=draw_thickness,
            )

        for start, end in self.__mask_list:
            # 将掩码区域涂黑
            image = cv2.rectangle(
                image,
                (int(start.x), int(start.y)),
                (int(end.x), int(end.y)),
                color=(0, 0, 0),
                thickness=draw_thickness,
            )

        # 标出对焦船头点
        if len(self.__zoom_list) != 0:
            image = cv2.circle(
                image,
                (int(self.__zoom_point.x), int(self.__zoom_point.y)),
                radius=10,
                color=self.__color,
                thickness=-1,
            )

        # 将运动船只画出来,白框
        # for sports_ship in self.__sports_ship_list:
        #     image = cv2.rectangle(
        #         image,
        #         (int(sports_ship.ship_start.x), int(sports_ship.ship_start.y)),
        #         (int(sports_ship.ship_end.x), int(sports_ship.ship_end.y)),
        #         color=(255, 255, 255),
        #         thickness=draw_thickness,
        #     )

        data[self.__out_key] = image
        return data

    def __output_image(self, data: Dict) -> Dict:
        image = data[self.__in_key][0].copy()
        draw_thickness = min(image.shape[:2]) // 320

        for ship_data in self.__ship_data_list:
            # 画出追踪 ID
            cv2.putText(
                image,
                str(ship_data.track_id),
                (int(ship_data.ship_start.x), int(ship_data.ship_start.y)),
                cv2.FONT_HERSHEY_COMPLEX,
                1,
                self.__color,
                2,
            )
            # 画出追踪框
            image = cv2.rectangle(
                image,
                (int(ship_data.ship_start.x), int(ship_data.ship_start.y)),
                (int(ship_data.ship_end.x), int(ship_data.ship_end.y)),
                color=self.__color,
                thickness=draw_thickness,
            )

        data[self.__out_key] = image
        return data

    def __calculate_offset(self, A: ShipData, B: ShipData) -> float:
        """
        计算两个坐标之间的横坐标距离
        """
        return A.get_midpoint_data().x - B.get_midpoint_data().x

    def __add_ship_data(self):
        for ship_data in self.__ship_data_list:
            # 找老版本的数据
            old_ship_data = self.__ship_processing_map.get(
                ship_data.track_id, ship_data
            )

            if ship_data.track_id not in self.__position_offset_map:
                self.__position_offset_map[ship_data.track_id] = deque()
            self.__position_offset_map[ship_data.track_id].append(
                self.__calculate_offset(old_ship_data, ship_data)
            )

            self.__ship_processing_map[ship_data.track_id] = ship_data

    def __get_sports_ship_list(self, position_offset_threshold):
        list_result = []

        # 过了阈值的,认为船移动
        for key in self.__ship_processing_map:
            position_offset_list = self.__position_offset_map[key]

            # position_offset_list 的 deque 队列求平均误差
            position_offset = sum(position_offset_list) / len(position_offset_list)

            # 偏差可能为正负数,以此区分航向
            if abs(position_offset) > position_offset_threshold:
                if self.__ship_heading == Heading.kLeft and position_offset > 0:
                    list_result.append(self.__ship_processing_map[key])
                elif self.__ship_heading == Heading.kRight and position_offset < 0:
                    list_result.append(self.__ship_processing_map[key])

        # 根据航向对船只进行排序:
        if self.__ship_heading == Heading.kLeft:
            # 从上到下,从左到右
            list_result = sorted(
                list_result, key=lambda box: (box.ship_start.y, box.ship_start.x)
            )
        elif self.__ship_heading == Heading.kRight:
            # 从上到下,从右到左
            list_result = sorted(
                list_result, key=lambda box: (box.ship_start.y, -box.ship_start.x)
            )

        # 清空所有缓存
        self.__ship_processing_map.clear()
        self.__position_offset_map.clear()

        return list_result

    def __detect_ready(self):
        if len(self.__ship_data_list) == 0:
            logging.info(f"当前未检测到船只,进行等待...")
            return StateMachineEnum.kDetectReady, 2

        self.__detect_count += 1
        # 积累数据
        self.__add_ship_data()

        if self.__detect_count < self.__detect_threshold:
            return StateMachineEnum.kDetectReady, 0

        # 积累数据后,获取运动船只
        self.__detect_count = 0
        self.__sports_ship_list = self.__get_sports_ship_list(
            self.__position_offset_threshold * 2
        )

        if len(self.__sports_ship_list) == 0:
            logging.info(f"当前未检测到运动船只,进行等待...")
            return StateMachineEnum.kDetectReady, 2

        logging.info("发现运动船只,航线暂停...")
        self.__http_message.flighttask_pause()
        self.__http_message.flight_authority_grab()
        self.__http_message.payload_authority_grab()
        self.__http_message.live_lens_change(CameraType.kWide)
        return StateMachineEnum.kDetect, 2

    def __detect(self):
        self.__detect_count += 1
        # 积累数据
        self.__add_ship_data()

        if self.__detect_count < self.__detect_threshold:
            return StateMachineEnum.kDetect, 0

        # 积累数据后,获取运动船只
        self.__detect_count = 0
        self.__sports_ship_list = self.__get_sports_ship_list(
            self.__position_offset_threshold
        )

        for ship_data in self.__sports_ship_list:
            y = (ship_data.ship_end.y + ship_data.ship_start.y) / 2

            # 变焦位置根据船只航向进行偏移
            if self.__ship_heading == Heading.kLeft:
                x = ship_data.ship_start.x

            elif self.__ship_heading == Heading.kRight:
                x = ship_data.ship_end.x

            target_width = self.__screen_size_x / (
                ship_data.ship_end.x - ship_data.ship_start.x
            )
            target_height = self.__screen_size_y / (
                ship_data.ship_end.y - ship_data.ship_start.y
            )
            # 计算变焦倍率
            self.__zoom_factor = max(target_width, target_height)
            # 缩小放大倍率是为了让画面尽可能有整艘船只
            self.__zoom_factor = int(self.__zoom_factor - 1)

            logging.info(
                f"找到船头点 : ({x} , {y}) -> (0.5 , {y / self.__screen_size_y}),缩放倍率 : {self.__zoom_factor}"
            )

            self.__http_message.camera_aim(
                CameraType.kWide,
                0.5,
                y / self.__screen_size_y,
            )

            # 记录检测历史
            self.__zoom_list.append(
                (
                    PointData(ship_data.ship_start.x, ship_data.ship_start.y),
                    PointData(ship_data.ship_end.x, ship_data.ship_end.y),
                )
            )
            self.__zoom_point = PointData(x, y)

            return StateMachineEnum.kZoom, 2

        # 所有船只检测完毕
        logging.info("检测结束,航线恢复...")
        self.__zoom_list.clear()
        self.__http_message.flighttask_recovery()
        return StateMachineEnum.kDetectReady, 10

    def __zoom(self):
        logging.info("切换到变焦镜头,镜头缩放...")
        self.__http_message.camera_focal_length_set(
            CameraType.kZoom, self.__zoom_factor
        )
        self.__http_message.live_lens_change(CameraType.kZoom)
        return StateMachineEnum.kZoomDetect, 2

    def __zoom_detect(self):
        self.__mask_list.clear()

        # 等待船只入画
        if len(self.__ship_data_list) == 0:
            logging.info(f"变焦未检测到船只,进行等待...")
            self.__zoom_sleep_count += 1

        zoom_max_area = 0

        for ship_data in self.__ship_data_list:
            # zoom_max_area = max(zoom_max_area, ship_data.get_area())
            self.__mask_list.append((ship_data.ship_start, ship_data.ship_end))

        # 中间有检测到一次就继续保持数秒
        if len(self.__mask_list) != 0:
            self.__zoom_sleep_count = 0

        # 检测框缩小到最大检测框的 1/3,退出检测
        # if self.__zoom_max_area != 0 and self.__zoom_max_area / 3 > zoom_max_area:
        #     logging.info(f"检测框缩小到最大检测框的 1/3,退出检测...")
        #     self.__zoom_sleep_count = self.__zoom_sleep_threshold

        if self.__zoom_sleep_count < self.__zoom_sleep_threshold:
            # self.__zoom_max_area = max(self.__zoom_max_area, zoom_max_area)
            return StateMachineEnum.kZoomDetect, 0

        # 等待镜头还原
        logging.info("等待设备恢复...")
        self.__zoom_sleep_count = 0
        # self.__zoom_max_area = 0
        self.__mask_list.clear()
        return StateMachineEnum.kCameraReset, 2

    def __camera_reset(self):
        # 广角镜头切换完毕,继续检测
        logging.info("切换到广角镜头,等待镜头恢复...")
        self.__http_message.live_lens_change(CameraType.kWide)
        self.__http_message.get_recovery_camera()
        return StateMachineEnum.kDetect, 2

    def _call(self, data: Dict) -> Dict:
        # 读取完毕关闭任务
        if data.get(self.__finish_tag_key, False):
            logging.info("任务完成")
            self.close()

        # 计算出调用该函数的帧率
        # fps = 1 / (time.time() - self.__last_fps_time)
        # logging.info(f"帧率 : {fps}")
        # self.__last_fps_time = time.time()

        # 获取船只数据
        data = self.__refresh_ship_data(data)

        # 运行状态机
        self.__state_machine.run_once()

        # 画出状态机流程的检测框
        data = self.__detect_image(data)

        # 画出所有检测船只
        # data = self.__output_image(data)

        return data

    def _init_check(self):
        """
        任务执行前初始化检查参数
        """

        if self.__debug:
            mkv_path = os.path.join(self.__image_output_path, "output.mkv")
            self.__out = cv2.VideoWriter(
                mkv_path,
                cv2.VideoWriter_fourcc(*"XVID"),
                30,
                (self.__screen_size_x, self.__screen_size_y),
            )

    def _set(self, data: Dict):
        if data.get("Left_Heading", False):
            self.__set_ship_heading(data["Left_Heading"])

    def _start(self):
        self._init_check()

    def _close(self):
        if self.__debug:
            self.__out.release()

        if not self.__test:
            esdk.StopH264Stream()


if __name__ == "__main__":
    pass
